package systemf;

import java.util.HashMap;
import java.util.Map;

/**
 * System F类型检查器
 * 实现类型推断和类型检查
 */
public class TypeChecker {
    
    private final Map<String, Type> context;
    
    public TypeChecker() {
        this.context = new HashMap<>();
    }
    
    public TypeChecker(Map<String, Type> context) {
        this.context = new HashMap<>(context);
    }
    
    /**
     * 类型检查：检查项M是否具有类型τ
     */
    public boolean checkType(Term term, Type type) {
        Type inferred = inferType(term);
        return inferred != null && typeEquals(inferred, type);
    }
    
    /**
     * 类型推断：推断项M的类型
     */
    public Type inferType(Term term) {
        if (term instanceof Term.Var) {
            Term.Var var = (Term.Var) term;
            Type type = context.get(var.getName());
            if (type == null) {
                throw new TypeException("未定义的变量: " + var.getName());
            }
            return type;
        } else if (term instanceof Term.Lambda) {
            Term.Lambda lambda = (Term.Lambda) term;
            TypeChecker newChecker = new TypeChecker(context);
            newChecker.context.put(lambda.getVar(), lambda.getType());
            Type bodyType = newChecker.inferType(lambda.getBody());
            return new Type.ArrowType(lambda.getType(), bodyType);
        } else if (term instanceof Term.App) {
            Term.App app = (Term.App) term;
            Type funcType = inferType(app.getFunction());
            if (!(funcType instanceof Type.ArrowType)) {
                throw new TypeException("应用左侧必须是函数类型，得到: " + funcType);
            }
            Type.ArrowType arrowType = (Type.ArrowType) funcType;
            Type argType = inferType(app.getArgument());
            if (!typeEquals(arrowType.getDomain(), argType)) {
                throw new TypeException("参数类型不匹配: 期望 " + arrowType.getDomain() + 
                                       ", 得到 " + argType);
            }
            return arrowType.getCodomain();
        } else if (term instanceof Term.TypeAbs) {
            Term.TypeAbs typeAbs = (Term.TypeAbs) term;
            TypeChecker newChecker = new TypeChecker(context);
            // 类型变量在类型环境中不需要存储，因为类型抽象是类型级别的
            Type bodyType = newChecker.inferType(typeAbs.getBody());
            return new Type.ForallType(typeAbs.getTypeVar(), bodyType);
        } else if (term instanceof Term.TypeApp) {
            Term.TypeApp typeApp = (Term.TypeApp) term;
            Type termType = inferType(typeApp.getTerm());
            if (!(termType instanceof Type.ForallType)) {
                throw new TypeException("类型应用左侧必须是全称类型，得到: " + termType);
            }
            Type.ForallType forallType = (Type.ForallType) termType;
            return forallType.getBody().substitute(forallType.getTypeVar(), typeApp.getType());
        }
        throw new TypeException("未知的项类型: " + term.getClass());
    }
    
    /**
     * 类型相等性检查（考虑α-等价）
     */
    private boolean typeEquals(Type t1, Type t2) {
        if (t1 == t2) return true;
        if (t1 == null || t2 == null) return false;
        
        if (t1 instanceof Type.BaseType && t2 instanceof Type.BaseType) {
            return ((Type.BaseType) t1).getName().equals(((Type.BaseType) t2).getName());
        }
        
        if (t1 instanceof Type.TypeVar && t2 instanceof Type.TypeVar) {
            return ((Type.TypeVar) t1).getName().equals(((Type.TypeVar) t2).getName());
        }
        
        if (t1 instanceof Type.ArrowType && t2 instanceof Type.ArrowType) {
            Type.ArrowType a1 = (Type.ArrowType) t1;
            Type.ArrowType a2 = (Type.ArrowType) t2;
            return typeEquals(a1.getDomain(), a2.getDomain()) &&
                   typeEquals(a1.getCodomain(), a2.getCodomain());
        }
        
        if (t1 instanceof Type.ForallType && t2 instanceof Type.ForallType) {
            Type.ForallType f1 = (Type.ForallType) t1;
            Type.ForallType f2 = (Type.ForallType) t2;
            // 简单的相等性检查（完整的α-等价需要更复杂的实现）
            return f1.getTypeVar().equals(f2.getTypeVar()) &&
                   typeEquals(f1.getBody(), f2.getBody());
        }
        
        return false;
    }
    
    /**
     * 添加变量到上下文
     */
    public void addVar(String name, Type type) {
        context.put(name, type);
    }
    
    /**
     * 类型异常
     */
    public static class TypeException extends RuntimeException {
        public TypeException(String message) {
            super(message);
        }
    }
}

